#include "socks.h"
#include <errno.h>
#include <string.h>
#include "client.h"
#include "log.h"
#include "config.h"
#include "fifo.h"
#include "misc.h"
#include "io.h"
#include "crypt.h"
#include "resolv.h"

#define SOCKS4_BASE_LEN     (1 + 1 + 2 + 4)
#define SOCKS5_BASE_LEN     (1 + 1 + 1 + 1 + 4 + 2)
#define SOCKS5_REQ_MAX_LEN  (1 + 1 + 1 + 1 + 1+255 + 2)

enum socks_cmd {
    CMD_CONNECT = 1,
    CMD_BIND = 2,
    CMD_UDP_ASSOCIATE = 3,
};

enum socks4_rep_code {
    S4_REP_GRANT    = 90, // request granted
    S4_REP_FAILED   = 91, // request rejected or failed
    S4_REP_CANNOT_CONN  = 92, // cannot connect to identd on the client
    S4_REP_REJECT   = 93, // the client program and identd report different user-ids
};

enum socks5_atyp {
    S5_ATYP_IPV4 = 0x01,
    S5_ATYP_DOMAINNAME = 0x03,
    S5_ATYP_IPV6 = 0x04
};

enum socks5_rep_code {
    S5_REP_SUCCESSED = 0x00,
    S5_REP_GENERAL_FAILURE = 0x01,
    S5_REP_NOT_ALLOW = 0x02,
    S5_REP_NET_UNREACH = 0x03,
    S5_REP_HOST_UNREACH = 0x04,
    S5_REP_CONN_REFUSED = 0x05,
    S5_REP_TTL_EXPIRED = 0x06,
    S5_REP_CMD_NOT_SUPPORT = 0x07,
    S5_REP_ATYP_NOT_SUPPORT = 0x08,
};

static void server_addr_resolv_cb(const char *domain, const struct addr *result,
        void *udata);

static struct request *socks_req_info_realloc(struct request *old) {
    struct request *new;

    if (!(new = request_realloc(old))) {
        // XXX: no mem ?
        free(old);
        return NULL;
    }

    if (new->socks.s4a_domain)
        new->socks.s4a_domain += (void *) new - (void *) old;

    return new;
}

static int socks4_reply(struct client *c, enum socks4_rep_code rep) {
    c->req->buf.socks.s4.vn = 0;
    c->req->buf.socks.s4.cd = rep;
    c->req->buf.len = SOCKS4_BASE_LEN;
    if (client_write2sock(c->lsock, c->req->buf.raw, c->req->buf.len, c) < 0)
        return -1;
    return fifo_is_empty(c->lout) ? 0 : 1;
}

static int socks5_reply(struct client *c, enum socks5_rep_code rep) {
    c->req->buf.socks.s5.rep = rep;
    if (c->req->buf.len < SOCKS5_BASE_LEN) {
        c->req->buf.len = SOCKS5_BASE_LEN;
        c->req->buf.socks.s5.atyp = S5_ATYP_IPV4;
    }
    if (client_write2sock(c->lsock, c->req->buf.raw, c->req->buf.len, c) < 0)
        return -1;
    return fifo_is_empty(c->lout) ? 0 : 1;
}

static int reply(struct client *c, int rep) {
    if (c->req->buf.socks.ver == SOCKS5)
        return socks5_reply(c, rep);
    else
        return socks4_reply(c, rep);
}

/* convert errno to SOCKS4/SOCKS5 reply code */
static int connect_failed(struct client *c) {
    int rep = c->req->buf.socks.ver == SOCKS5 ? S5_REP_GENERAL_FAILURE : S4_REP_FAILED;
    int level = LOG_INFO;

    switch(errno) {
    case ECONNREFUSED:
        rep = c->req->buf.socks.ver == SOCKS5 ? S5_REP_CONN_REFUSED : S4_REP_CANNOT_CONN;
        break;
    case ENETUNREACH:
        rep = c->req->buf.socks.ver == SOCKS5 ? S5_REP_NET_UNREACH : S4_REP_CANNOT_CONN;
        break;
    case EHOSTUNREACH:
        rep = c->req->buf.socks.ver == SOCKS5 ? S5_REP_HOST_UNREACH : S4_REP_CANNOT_CONN;
        break;
    case ETIMEDOUT:
        break;
    default:
        level = LOG_WARNING;
        break;
    }

    client_log_socket_error(level, c, c->rsock, "connect failure");

    update_server_weight(c, -1);

    return rep;
}

static void remote_connected_handler(EV_P_ ev_io *w, int revents) {
    ssize_t n;
    struct client *c = w->data;

    if ((errno = get_socket_errno(w->fd))) {
        reply(c, connect_failed(c));
        goto DESTROY;
    }

    ev_io_stop(EV_A_ w);

    update_server_weight(c, 1);

    if (crypt_init(&c->crypt, c->srv->crypt_method, c->srv->passwd))
        goto REPLY_FAILED;

    /* +------+----------+----------+
     * | ATYP | DST.ADDR | DST.PORT |
     * +------+----------+----------+
     * |  1   | Variable |    2     |
     * +------+----------+----------+ */
    if (c->req->buf.socks.ver == SOCKS5) {
        n = crypt_en(&c->crypt, c->ctx->outbuf, c->ctx->outbuflen,
                    c->req->buf.raw + 3, c->req->buf.len - 3);
    } else {
        unsigned char *p = c->ctx->inbuf;
        if (c->req->socks.s4a_domain) {
            p[0] = S5_ATYP_DOMAINNAME;
            p[1] = strlen(c->req->socks.s4a_domain);
            memcpy(p + 2, c->req->socks.s4a_domain, p[1]);
            *(in_port_t *) (p + 2 + p[1]) = c->req->buf.socks.s4.port;
            n = 2 + p[1] + 2;
        } else {
            p[0] = S5_ATYP_IPV4;
            *(in_addr_t *) (p + 1) = c->req->buf.socks.s4.ip;
            *(in_port_t *) (p + 1 + 4) = c->req->buf.socks.s4.port;
            n = 1 + 4 + 2;
        }
        n = crypt_en(&c->crypt, c->ctx->outbuf, c->ctx->outbuflen, c->ctx->inbuf, n);
    }
    if (n < 0)
        goto REPLY_FAILED;

    if (client_write2sock(w->fd, c->ctx->outbuf, n, c) < 0)
        goto REPLY_FAILED;

    if (reply(c, c->req->buf.socks.ver == SOCKS5 ? S5_REP_SUCCESSED : S4_REP_GRANT) < 0)
        goto DESTROY;

    ev_io_set(w, w->fd, EV_READ);
    w->data = c;
    ev_set_cb(w, relay_handler);
    ev_io_start(EV_A_ w);
    c->levio.data = c;
    ev_set_cb(&c->levio, relay_handler);
    ev_timer_stop(EV_A_ &c->evtimer);
    ev_timer_set(&c->evtimer, 0, c->ctx->cfg->socket_timeout);
    ev_set_cb(&c->evtimer, sweep_handler);
    ev_timer_again(EV_A_ &c->evtimer);

    return;

REPLY_FAILED:
    if (c->req->buf.socks.ver == SOCKS5)
        socks5_reply(c, S5_REP_GENERAL_FAILURE);
    else
        socks4_reply(c, S4_REP_FAILED);
DESTROY:
    client_destroy(c);
}

static void connect_timer_cb(EV_P_ ev_timer *watcher, int revents) {
    struct client *c = watcher->data;
    errno = ETIMEDOUT;
    reply(c, connect_failed(c));
    client_destroy(c);
}

static void try_connect(struct client *c, const struct addr *addr) {
    union {
        struct sockaddr_in in4;
        struct sockaddr_in6 in6;
    } in;
    socklen_t inlen;
    struct sockaddr *inp = (struct sockaddr *)&in;
    int repcode = c->req->buf.socks.ver == SOCKS5 ? S5_REP_GENERAL_FAILURE : S4_REP_FAILED;

    memset(&in, 0, sizeof(in));
    if (addr->addr[0].family == AF_INET) {
        in.in4.sin_family = AF_INET;
        in.in4.sin_addr = addr->addr[0].in4;
        in.in4.sin_port = c->srv->port;
        inlen = sizeof(in.in4);
    } else {
        in.in6.sin6_family = AF_INET6;
        in.in6.sin6_addr = addr->addr[0].in6;
        in.in6.sin6_port = c->srv->port;
        inlen = sizeof(in.in6);
    }

    sockaddr_ntop(inp, c->raddrinfo, sizeof(c->raddrinfo));

    if ((c->rsock = socket(inp->sa_family, SOCK_STREAM, 0)) < 0) {
        LOGW("Cannot create socket for %s(%s): %s", c->raddrinfo, c->srv->addr, strerror(errno));
        goto DESTROY;
    }
    nonblocking(c->rsock);
    if (safe_connect(c->rsock, inp, inlen) < 0 && errno != EINPROGRESS) {
        repcode = connect_failed(c);
        goto DESTROY;
    }

    ev_io_init(&c->revio, remote_connected_handler, c->rsock, EV_WRITE);
    ev_io_start(c->ctx->evloop, &c->revio);

    ev_timer_set(&c->evtimer, 0., c->ctx->cfg->connect_timeout);
    ev_set_cb(&c->evtimer, connect_timer_cb);
    ev_timer_again(c->ctx->evloop, &c->evtimer);

    return;

DESTROY:
    if (c->req->buf.socks.ver == SOCKS5)
        socks5_reply(c, repcode);
    else
        socks4_reply(c, repcode);
    client_destroy(c);
}

static void server_addr_resolv_cb(const char *domain, const struct addr *result,
        void *udata) {
    struct client *c = udata;

    c->resolv = NULL;
    c->resolv_cb = NULL;

    if (result && result->cnt) {
        try_connect(c, result);
    } else {
        update_server_weight(c, -1);
        if (c->req->buf.socks.ver == SOCKS5)
            socks5_reply(c, S5_REP_GENERAL_FAILURE);
        else
            socks4_reply(c, S4_REP_FAILED);
        client_destroy(c);
    }
}

static void request_addr2str(const struct request *req, char *strbuf, size_t buflen) {
    union {
        struct sockaddr_in in;
        struct sockaddr_in6 in6;
    } addr;

    switch (req->buf.socks.ver) {
    case SOCKS4:
        if (req->socks.s4a_domain) {
            size_t len = strlen(req->socks.s4a_domain);
            memcpy(strbuf, req->socks.s4a_domain, len);
            snprintf(strbuf + len, buflen - len, ":%hu", ntohs(req->buf.socks.s4.port));
            return;
        } else {
            addr.in.sin_family = AF_INET;
            addr.in.sin_addr.s_addr = req->buf.socks.s4.ip;
            addr.in.sin_port = req->buf.socks.s4.port;
            break;
        }

    case SOCKS5:
        switch (req->buf.socks.s5.atyp) {
        case S5_ATYP_IPV4:
            addr.in.sin_family = AF_INET;
            addr.in.sin_addr = req->buf.socks.s5.in4.addr;
            addr.in.sin_port = req->buf.socks.s5.in4.port;
            break;
        case S5_ATYP_IPV6:
            addr.in6.sin6_family = AF_INET6;
            addr.in6.sin6_addr = req->buf.socks.s5.in6.addr;
            addr.in6.sin6_port = req->buf.socks.s5.in6.port;
            break;
        case S5_ATYP_DOMAINNAME:
            memcpy(strbuf, req->buf.socks.s5.domain.domain, req->buf.socks.s5.domain.len);
            in_port_t *port = (in_port_t *)(req->buf.socks.s5.domain.domain
                    + req->buf.socks.s5.domain.len);
            snprintf(strbuf + req->buf.socks.s5.domain.len,
                    buflen - req->buf.socks.s5.domain.len, ":%hu", ntohs(*port));
            return;
        default:
            addr.in.sin_family = AF_UNSPEC;
            break;
        }
        break;
    default:
        return;
    }
    sockaddr_ntop(&addr, strbuf, buflen);
}

/* return -1 if client destroyed */
static int socks4_connect(struct client *c, struct request *req) {
    if (req->buf.len <= SOCKS4_BASE_LEN || req->buf.raw[req->buf.len - 1] != 0)
        return 0;

    unsigned char *ip = (unsigned char *) &req->buf.socks.s4.ip;
    if (!ip[0] && !ip[1] && !ip[2] && ip[3]) {
        /* is socks4a */
        if (!req->socks.s4a_domain) {
            char *p = req->buf.raw + SOCKS4_BASE_LEN;
            while (*p++) ;
            req->socks.s4a_domain = (char *)p;
            if (p > req->buf.raw + req->buf.len)
                return 0; /* we need more data for domain */
        }
    }

    if (req->socks.s4a_domain && strlen(req->socks.s4a_domain) > 255) {
        LOGW("%s domain name too long \"%s\" (SOCKS4a)", c->laddrinfo, req->socks.s4a_domain);
        socks4_reply(c, S4_REP_FAILED);
        goto DESTROY;
    }

    c->src_proto = req->socks.s4a_domain ? SP_SOCKS4a : SP_SOCKS4;
    request_addr2str(req, c->dst_info, sizeof(c->dst_info));

    LOGI("%s %s --> %s --> %s", src_proto_str(c->src_proto), c->laddrinfo, c->srv->addr, c->dst_info);

    ev_timer_stop(c->ctx->evloop, &c->evtimer);

    /* resolve socks server domain */
    c->resolv_cb = server_addr_resolv_cb;
    c->resolv = resolv_query(c->ctx->resolver, c->srv->addr, server_addr_resolv_cb, c);
    if (!c->resolv && errno) {
        c->resolv_cb = NULL;
        socks4_reply(c, S4_REP_FAILED);
        goto DESTROY;
    }

    return 0;

DESTROY:
    client_destroy(c);
    return -1;
}

/* return -1 if client destroyed */
static int socks4_addr_read_event(int fd, struct client *c) {
    ssize_t n;

    client_update_activity(c);

    /* we not need recv() if have fresh data */
    if (c->req->buf.pos != c->req->buf.raw) {
        if (c->req->buf.len >= c->req->buf.size) {
            if (c->req->buf.size >= c->ctx->cfg->buf_size) {
                LOGW("%s recv socks4 request error: Request too long", c->laddrinfo);
                goto DESTROY;
            }
            if (!(c->req = socks_req_info_realloc(c->req))) {
                LOGW("%s recv socks4 request error: Not enough memory", c->laddrinfo);
                socks4_reply(c, S4_REP_FAILED);
                goto DESTROY;
            }
        }

        if ((n = read(fd, c->req->buf.raw + c->req->buf.len,
                c->req->buf.size - c->req->buf.len)) < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
                return 0;
            client_log_socket_error(LOG_WARNING, c, fd, "recv error");
            goto DESTROY;
        }
        if (!n)
            goto DESTROY;
        c->req->buf.len += n;
    }

    ++c->req->buf.pos;

    if (c->req->buf.len < 2)
        return 0;

    switch (c->req->buf.socks.s4.cd) {
    case CMD_CONNECT:
        return socks4_connect(c, c->req);
    case CMD_BIND:
        // TODO: return socks4_bind();
    default:
        LOGW("%s unsupported SOCKS4 command (%02hhX)", c->laddrinfo, c->req->buf.socks.s4.cd);
        socks4_reply(c, S4_REP_FAILED);
        break;
    }

DESTROY:
    client_destroy(c);
    return -1;
}

void socks4_addr_handler(EV_P_ ev_io *w, int revents) {
    struct client *c = w->data;

    if ((revents & EV_READ) && socks4_addr_read_event(w->fd, c))
        return;

    if (revents & EV_WRITE)
        client_write_event(EV_A_ w, c);
}

/* return -1 if client destroyed */
static int socks5_connect(struct client *c, struct request *req) {
    /* check RSV and ATYP */
    if (req->buf.len < 4)
        return 0; /* need more */
    size_t dataLen = 6;
    switch (req->buf.socks.s5.atyp) {
    case S5_ATYP_IPV4:
        dataLen += 4;
        break;
    case S5_ATYP_IPV6:
        dataLen += 16;
        break;
    case S5_ATYP_DOMAINNAME:
        if (req->buf.len < 5)
            return 0; /* need more */
        dataLen += 1 + req->buf.socks.s5.domain.len;
        break;
    default:
        LOGW("%s unsupported SOCKS5 ATYPE (%02hhX)", c->laddrinfo, req->buf.socks.s5.atyp);
        socks5_reply(c, S5_REP_ATYP_NOT_SUPPORT);
        goto DESTROY;
    }
    if (req->buf.len < dataLen)
        return 0; /* need more */

    c->src_proto = SP_SOCKS5;
    request_addr2str(req, c->dst_info, sizeof(c->dst_info));

    LOGI("%s %s --> %s --> %s", src_proto_str(c->src_proto), c->laddrinfo, c->srv->addr, c->dst_info);

    ev_timer_stop(c->ctx->evloop, &c->evtimer);

    /* resolve socks server domain */
    c->resolv_cb = server_addr_resolv_cb;
    c->resolv = resolv_query(c->ctx->resolver, c->srv->addr, server_addr_resolv_cb, c);
    if (!c->resolv && errno) {
        c->resolv_cb = NULL;
        socks5_reply(c, S5_REP_GENERAL_FAILURE);
        goto DESTROY;
    }

    return 0;

DESTROY:
    client_destroy(c);
    return -1;
}

/* return -1 if client destroyed */
static int socks5_addr_read_event(int fd, struct client *c) {
    ssize_t n;

    client_update_activity(c);

    if ((n = read(fd, c->req->buf.raw + c->req->buf.len,
            c->req->buf.size - c->req->buf.len)) < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
            return 0;
        LOGW("%s recv error: %s", c->laddrinfo, strerror(errno));
        goto DESTROY;
    }
    if (!n)
        goto DESTROY;
    c->req->buf.len += n;

    /* check and process */

    if (c->req->buf.socks.ver != SOCKS5) {
        LOGW("%s invalid SOCKS version code %02hhX", c->laddrinfo, c->req->buf.socks.ver);
        goto DESTROY;
    }
    /* check CMD */
    if (c->req->buf.len < 2)
        return 0; /* need more */

    switch (c->req->buf.socks.s5.cmd) {
    case CMD_CONNECT:
        return socks5_connect(c, c->req);
    case CMD_BIND:
        // TODO:
    case CMD_UDP_ASSOCIATE:
        // TODO:
    default:
        LOGW("%s unsupported SOCKS5 command (%02hhX)", c->laddrinfo, c->req->buf.socks.s5.cmd);
        socks5_reply(c, S5_REP_CMD_NOT_SUPPORT);
        goto DESTROY;
    }

DESTROY:
    client_destroy(c);
    return -1;
}

static void socks5_addr_handler(EV_P_ ev_io *w, int revents) {
    struct client *c = w->data;

    if ((revents & EV_READ) && socks5_addr_read_event(w->fd, c))
        return;

    if (revents & EV_WRITE)
        client_write_event(EV_A_ w, c);
}

/* return -1 if client destroyed */
static int hello_read_event(ev_io *w, struct client *c) {
    ssize_t n;

    client_update_activity(c);

    /* we not need recv() if have fresh data */
    if (c->req->buf.pos != c->req->buf.raw) {
        if ((n = read(w->fd, c->req->buf.raw + c->req->buf.len,
                c->req->buf.size - c->req->buf.len)) < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
                return 0;
            client_log_socket_error(LOG_WARNING, c, w->fd, "recv error");
            goto DESTROY;
        }
        if (!n)
            goto DESTROY;
        c->req->buf.len += n;
    }

    /* check */
    if (c->req->buf.socks.ver != SOCKS5) {
        LOGW("%s invalid SOCKS5 version code %d", c->laddrinfo, c->req->buf.socks.ver);
        goto DESTROY;
    }

    ++c->req->buf.pos;

    if (c->req->buf.len < 3)
        return 0;
    if (c->req->buf.len < 2UL + c->req->buf.raw[1])
        return 0;

    c->req->buf.raw[1] = 0x00;
    if (client_write2sock(c->lsock, c->req->buf.raw, 2, c) < 0)
        goto DESTROY;

    c->req->buf.len = 0;
    c->req->buf.pos = c->req->buf.raw;
    ev_set_cb(w, socks5_addr_handler);

    return 0;

DESTROY:
    client_destroy(c);
    return -1;
}

void socks5_hello_handler(EV_P_ ev_io *w, int revents) {
    struct client *c = w->data;

    if ((revents & EV_READ) && hello_read_event(w, c))
        return;

    if (revents & EV_WRITE)
        client_write_event(EV_A_ w, c);
}
